import os
import glob
from subprocess import call


def extract_dataset_frames(data_dir):
    if not os.path.exists(os.path.join(data_dir, 'images')):
        os.makedirs(os.path.join(data_dir, 'images'))

    for filepath in glob.glob(os.path.join(data_dir, 'videos/*/*/*.avi')):
        splitted_path = filepath.split(os.path.sep)
        data_set = splitted_path[-3]
        category = splitted_path[-2]
        filename = splitted_path[-1]

        dest = os.path.join(data_dir, 'images', data_set, category, filename[:-4],
                            filename[:-4] + '-%04d.jpg')

        if not os.path.exists(os.path.join(data_dir, 'images', data_set, category, filename[:-4])):
            os.makedirs(os.path.join(data_dir, 'images', data_set, category, filename[:-4]))

        print('Extracting ' + category + '/' + filename)
        call(["ffmpeg", "-i", filepath, dest])


if __name__ == '__main__':
    download_dir = '/home/ww/CrowdFights'
    extract_dataset_frames(download_dir)
    # for split in ['01', '02', '03']:
    #     train_test_split(download_dir, split)
